import os
import inspect
import random
import datetime
import argparse

import torch
import xlsxwriter
import ray


class Function(list):
    def __call__(self, x):
        raise NotImplementedError()


class A(Function):
    def __init__(self, **kwargs):
        lower, upper = kwargs['lower'], kwargs['upper']
        super().__init__([random.randint(lower, upper) for _ in range(6)])

    def __call__(self, x):
        return (x + self[0]) * self[1] / self[2] + self[3] / self[4] * x - self[5]

    def __repr__(self):
        return f"(x + {self[0]}) * {self[1]} / {self[2]} + {self[3]} / {self[4]} * x = {self[5]}"


@ray.remote
class Worker(object):
    def __init__(self, functions, **kwargs):
        self.functions = functions
        self.kwargs = kwargs

    def __call__(self):
        kwargs = self.kwargs
        function = random.choice(self.functions)(**kwargs)
        last = None
        while True:
            x = torch.tensor(random.randint(kwargs['lower'], kwargs['upper']), dtype=torch.float, requires_grad=True)
            optimizer = torch.optim.SGD([x], lr=1)
            for _ in range(kwargs['step']):
                optimizer.zero_grad()
                y = function(x)
                if torch.isnan(y).item() or y.item() == 0:
                    break
                loss = torch.abs(y)
                loss.backward()
                grad = torch.sign(x.grad)
                if last is not None and grad != last:
                    break
                x.grad = last = grad
                optimizer.step()
            if y.item() == 0:
                return repr(function), x.item()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-r', '--rows', type=int, default=36)
    parser.add_argument('-s', '--size', type=int, default=16)
    parser.add_argument('-l', '--lower', type=int, default=-20)
    parser.add_argument('-u', '--upper', type=int, default=20)
    parser.add_argument('--step', type=int, default=100000)
    parser.add_argument('--address')
    parser.add_argument('--local_mode', action='store_true')
    args = parser.parse_args()
    ray.init(args.address, local_mode=args.local_mode)
    functions = [c for c in globals().values() if c != Function and inspect.isclass(c) and issubclass(c, Function)]
    pool = ray.util.ActorPool([Worker.remote(functions, **args.__dict__) for _ in range(1 if args.local_mode else int(ray.available_resources()['CPU']))])
    results = pool.map(lambda a, v: a.__call__.remote(), [None] * args.rows)
    root = os.path.splitext(__file__)[0]
    os.makedirs(root, exist_ok=True)
    with xlsxwriter.Workbook(os.path.join(root, datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S.%f') + '.xlsx')) as workbook:
        worksheet = workbook.add_worksheet('题目')
        worksheet_answer = workbook.add_worksheet('答案')
        format = workbook.add_format({'font_size': args.size})
        for row, (function, x) in enumerate(results):
            worksheet.write(row, 0, function, format)
            worksheet_answer.write(row, 0, f"{x}", format)


if __name__ == '__main__':
    main()